#ifndef NETS_HPP
#define NETS_HPP

#include "torch/torch.h"

class CustomDtaset : public torch::data::Dataset<CustomDtaset>
{
private:
    torch::Tensor inputs_; ///< 输入
    torch::Tensor labels_; ///< 标签

public:
    CustomDtaset(torch::Tensor input_data, torch::Tensor label_data);
    torch::data::Example<> get(size_t index) override;
    torch::optional<size_t> size() const override;
};

class Net : public torch::nn::Module
{
public:
    Net(const int inDim, const int outDim);

    at::Tensor forward(at::Tensor t);

private:
    torch::nn::Linear phase{nullptr};
    torch::nn::Linear trans{nullptr};
};

#endif // NETS_HPP
